import warnings
from collections import Counter

import numpy as np
import pandas as pd
from joblib import load
from rdkit import RDLogger, DataStructs
from rdkit.Chem import AllChem, MACCSkeys

import similaritycalc as sim

# ignore warnings
warnings.filterwarnings('ignore')
RDLogger.DisableLog('rdApp.*')

# load model
clf = load('clf.joblib')
mlp = load('mlp.joblib')

# import template file and test compounds
df = pd.read_excel("database_0714.xlsx")
print(df.head(1))
data = pd.read_excel("test set 715.xlsx")
data["product_list"] = None
data["product_dict_list"] = None
print(data.head(1))

# generate predicted products
for i in range(len(data)):
    try:
        reactant = data['substrate_smiles_canonical'][i]
        m = AllChem.MolFromSmiles(sim.canonicalize_smiles(reactant))
        pubchem_id = data['PUBCHEM CID'][i]
        test_df = df.drop(df[df['PUBCHEM CID'] == pubchem_id].index)
        maccs_keys = MACCSkeys.GenMACCSKeys(m)
        array1 = np.zeros((0,), dtype=np.float64)
        DataStructs.ConvertToNumpyArray(maccs_keys, array1)
        fps = AllChem.GetMorganFingerprintAsBitVect(m, 3, nBits=2048)
        array2 = np.zeros((0,), dtype=np.float64)
        DataStructs.ConvertToNumpyArray(fps, array2)
        array3 = np.append(array1, array2, axis=0)
        test_arr_2d = np.reshape(array3, (1, 2215))
        array = mlp.predict_proba(test_arr_2d)

        result_lst = []

        for j in range(len(test_df)):
            try:
                product_list = sim.apply_templates(reactant, test_df['extract templates'][j])
                score_r = sim.calculate_similarity(sim.get_fingerprint(reactant, "Morgan2"),
                                                   sim.get_fingerprint(test_df['substrate_smiles_canonical'][j],
                                                                       "Morgan2"),
                                                   "Tanimoto")
                reaction_class_pos = int(test_df['ML class no.'][j]) - 1
                # score_ml = array[reaction_class_pos][0][1]
                score_ml = array[0][reaction_class_pos]

                for product in product_list:
                    result_dict = {}
                    score_p = sim.calculate_similarity(sim.get_fingerprint(product, "Morgan2"),
                                                       sim.get_fingerprint(test_df['prod_smiles_canonical'][j],
                                                                           "Morgan2"),
                                                       "Tanimoto")
                    score = score_r * score_p + score_ml * 0.1  # could be adjusted
                    # score = score_r * score_p
                    result_dict['score'] = score
                    result_dict['product'] = product
                    result_dict['class'] = test_df['reaction class no.'][j]
                    result_dict['reaction_id'] = test_df['Reaction ID'][j]
                    result_lst.append(result_dict)
            except (Exception,):
                pass

        result_lst_sorted = sorted(result_lst, key=lambda a: a['score'], reverse=True)
        result_output_list = []
        result_dict_list = []
        for r in result_lst_sorted:
            if sim.canonicalize_smiles(r['product']) not in result_output_list:
                lst = []
                for res in result_dict_list:
                    lst.append(res['class'])
                counter = Counter(lst)
                if counter[r['class']] < 30:  # could be adjusted
                    result_output_list.append(sim.canonicalize_smiles(r['product']))
                    result_dict_list.append(r)
        print(i, "completed", result_output_list[:3])
        data['product_list'][i] = ';'.join(result_output_list)
        data['product_dict_list'][i] = result_dict_list
    except (Exception,):
        print(i, "error")

# get accuracy metrics
data["sim_ml_limited_class_precision"] = None
data["sim_ml_limited_class_recall"][i] = None

for i in range(len(data)):
    try:
        product_list = data['product sum'][i].split(";")
        sim_list = data['product_list'][i].split(";")
        data["sim_ml_limited_class_precision"][i] = sim.get_precison(3, sim_list, product_list)
        data["sim_ml_limited_class_recall"][i] = sim.get_recall(3, sim_list, product_list)
    except (Exception,):
        # print(i)
        data["sim_ml_limited_class_precision"][i] = 0
        data["sim_ml_limited_class_recall"][i] = 0

print(data["sim_ml_limited_class_precision"].mean(), data["sim_ml_limited_class_recall"].mean())
data.to_excel("test715.xlsx")
